import bpy
from ogreEx.CBone import Bone
from ogreEx.CReportSingleton import Report
from ogreEx.axis import swap
from ogreEx.config import getConfig
from ogreEx.xml.CRDocument import RDocument


# Skeleton
class Skeleton(object):
    def get_bone( self, name ):
        for b in self.bones:
            if b.name == name:
                 return b

    def __init__(self, ob ):
        if ob.location.x != 0 or ob.location.y != 0 or ob.location.z != 0:
            Report.warnings.append('ERROR: Mesh (%s): is offset from Armature - zero transform is required' %ob.name)
        if ob.scale.x != 1 or ob.scale.y != 1 or ob.scale.z != 1:
            Report.warnings.append('ERROR: Mesh (%s): has been scaled - scale(1,1,1) is required' %ob.name)

        self.object = ob
        self.bones = []
        mats = {}
        self.arm = arm = ob.find_armature()
        arm.hide = False
        self._restore_layers = list(arm.layers)
        #arm.layers = [True]*20      # can not have anything hidden - REQUIRED?

        for pbone in arm.pose.bones:
            if pbone.bone.use_deform or not getConfig()['ONLY_DEFORMABLE_BONES']:
                mybone = Bone( arm.data.bones[pbone.name] ,pbone, self )
                self.bones.append( mybone )

        if arm.name not in Report.armatures:
            Report.armatures.append( arm.name )

        ## bad idea - allowing rotation of armature, means vertices must also be rotated,
        ## also a bug with applying the rotation, the Z rotation is lost
        #x,y,z = arm.matrix_local.copy().inverted().to_euler()
        #e = mathutils.Euler( (x,z,y) )
        #self.object_space_transformation = e.to_matrix().to_4x4()
        x,y,z = arm.matrix_local.to_euler()
        if x != 0 or y != 0 or z != 0:
            Report.warnings.append('ERROR: Armature: %s is rotated - (rotation is ignored)' %arm.name)

        ## setup bones for Ogre format ##
        for b in self.bones: b.rebuild_tree()
        ## walk bones, convert them ##
        self.roots = []
        ep = 0.0001
        for b in self.bones:
            if not b.parent:
                b.compute_rest()
                loc,rot,scl = b.ogre_rest_matrix.decompose()
                #if loc.x or loc.y or loc.z:
                #    Report.warnings.append('ERROR: root bone has non-zero transform (location offset)')
                #if rot.w > ep or rot.x > ep or rot.y > ep or rot.z < 1.0-ep:
                #    Report.warnings.append('ERROR: root bone has non-zero transform (rotation offset)')
                self.roots.append( b )

    def to_xml( self ):
        _fps = float( bpy.context.scene.render.fps )

        doc = RDocument()
        root = doc.createElement('skeleton'); doc.appendChild( root )
        bones = doc.createElement('bones'); root.appendChild( bones )
        bh = doc.createElement('bonehierarchy'); root.appendChild( bh )
        for i,bone in enumerate(self.bones):
            b = doc.createElement('bone')
            b.setAttribute('name', bone.name)
            b.setAttribute('id', str(i) )
            bones.appendChild( b )
            mat = bone.ogre_rest_matrix.copy()
            if bone.parent:
                bp = doc.createElement('boneparent')
                bp.setAttribute('bone', bone.name)
                bp.setAttribute('parent', bone.parent.name)
                bh.appendChild( bp )

            pos = doc.createElement( 'position' ); b.appendChild( pos )
            x,y,z = mat.to_translation()
            pos.setAttribute('x', '%6f' %x )
            pos.setAttribute('y', '%6f' %y )
            pos.setAttribute('z', '%6f' %z )
            rot =  doc.createElement( 'rotation' ) # "rotation", not "rotate"
            b.appendChild( rot )

            q = mat.to_quaternion()
            rot.setAttribute('angle', '%6f' %q.angle )
            axis = doc.createElement('axis'); rot.appendChild( axis )
            x,y,z = q.axis
            axis.setAttribute('x', '%6f' %x )
            axis.setAttribute('y', '%6f' %y )
            axis.setAttribute('z', '%6f' %z )

            # Ogre bones do not have initial scaling? ##
            # note: Ogre bones by default do not pass down their scaling in animation,
            # so in blender all bones are like 'do-not-inherit-scaling'
            if 0:
                scale = doc.createElement('scale'); b.appendChild( scale )
                x,y,z = swap( mat.to_scale() )
                scale.setAttribute('x', str(x))
                scale.setAttribute('y', str(y))
                scale.setAttribute('z', str(z))

        arm = self.arm
        if not arm.animation_data or (arm.animation_data and not arm.animation_data.nla_tracks):  # assume animated via constraints and use blender timeline.
            anims = doc.createElement('animations'); root.appendChild( anims )
            anim = doc.createElement('animation'); anims.appendChild( anim )
            tracks = doc.createElement('tracks'); anim.appendChild( tracks )
            anim.setAttribute('name', 'my_animation')
            start = bpy.context.scene.frame_start; end = bpy.context.scene.frame_end
            anim.setAttribute('length', str( (end-start)/_fps ) )

            _keyframes = {}
            _bonenames_ = []
            for bone in arm.pose.bones:
                if self.get_bone(bone.name):                     #check if the bone was exported
                    _bonenames_.append( bone.name )
                    track = doc.createElement('track')
                    track.setAttribute('bone', bone.name)
                    tracks.appendChild( track )
                    keyframes = doc.createElement('keyframes')
                    track.appendChild( keyframes )
                    _keyframes[ bone.name ] = keyframes

            for frame in range( int(start), int(end)+1, bpy.context.scene.frame_step):
                bpy.context.scene.frame_set(frame)
                for bone in self.roots: bone.update()
                print('\t\t Frame:', frame)
                for bonename in _bonenames_:
                    if self.get_bone(bonename):                         #check if the bone was exported
                        bone = self.get_bone( bonename )
                        _loc = bone.pose_location
                        _rot = bone.pose_rotation
                        _scl = bone.pose_scale

                        keyframe = doc.createElement('keyframe')
                        keyframe.setAttribute('time', str((frame-start)/_fps))
                        _keyframes[ bonename ].appendChild( keyframe )
                        trans = doc.createElement('translate')
                        keyframe.appendChild( trans )
                        x,y,z = _loc
                        trans.setAttribute('x', '%6f' %x)
                        trans.setAttribute('y', '%6f' %y)
                        trans.setAttribute('z', '%6f' %z)

                        rot =  doc.createElement( 'rotate' )
                        keyframe.appendChild( rot )
                        q = _rot
                        rot.setAttribute('angle', '%6f' %q.angle )
                        axis = doc.createElement('axis'); rot.appendChild( axis )
                        x,y,z = q.axis
                        axis.setAttribute('x', '%6f' %x )
                        axis.setAttribute('y', '%6f' %y )
                        axis.setAttribute('z', '%6f' %z )

                        scale = doc.createElement('scale')
                        keyframe.appendChild( scale )
                        x,y,z = _scl
                        scale.setAttribute('x', '%6f' %x)
                        scale.setAttribute('y', '%6f' %y)
                        scale.setAttribute('z', '%6f' %z)


        elif arm.animation_data:
            anims = doc.createElement('animations'); root.appendChild( anims )
            if not len( arm.animation_data.nla_tracks ):
                Report.warnings.append('you must assign an NLA strip to armature (%s) that defines the start and end frames' %arm.name)

            if getConfig()['INDEPENDENT_ANIM']:
                for nla in arm.animation_data.nla_tracks:
                    nla.mute = True

            for nla in arm.animation_data.nla_tracks:        # NLA required, lone actions not supported
                if not len(nla.strips): print( 'skipping empty NLA track: %s' %nla.name ); continue
                print('NLA track:',  nla.name)

                if getConfig()['INDEPENDENT_ANIM']:
                    nla.mute = False
                    for strip in nla.strips:
                        strip.mute = True

                for strip in nla.strips:
                    if getConfig()['INDEPENDENT_ANIM']: strip.mute = False
                    print('   strip name:', strip.name)
                    anim = doc.createElement('animation'); anims.appendChild( anim )
                    tracks = doc.createElement('tracks'); anim.appendChild( tracks )
                    Report.armature_animations.append( '%s : %s [start frame=%s  end frame=%s]' %(arm.name, nla.name, strip.frame_start, strip.frame_end) )

                    anim.setAttribute('name', strip.name)                       # USE the strip name
                    anim.setAttribute('length', str( (strip.frame_end-strip.frame_start)/_fps ) )

                    stripbones = []
                    if getConfig()['ONLY_ANIMATED_BONES']:
                        for group in strip.action.groups:        # check if the user has keyed only some of the bones (for anim blending)
                            if group.name in arm.pose.bones: stripbones.append( group.name )
                        if not stripbones:                                    # otherwise we use all bones
                            stripbones = [ bone.name for bone in arm.pose.bones ]
                    else:
                        stripbones = [ bone.name for bone in arm.pose.bones ]

                    _keyframes = {}
                    for bonename in stripbones:
                        if self.get_bone(bonename):                     #check if the bone was exported
                            track = doc.createElement('track')
                            track.setAttribute('bone', bonename)
                            tracks.appendChild( track )
                            keyframes = doc.createElement('keyframes')
                            track.appendChild( keyframes )
                            _keyframes[ bonename ] = keyframes

                    for frame in range( int(strip.frame_start), int(strip.frame_end)+1, bpy.context.scene.frame_step):#thanks to Vesa
                        bpy.context.scene.frame_set(frame)
                        for bone in self.roots: bone.update()
                        for bonename in stripbones:
                            if self.get_bone( bonename ):               #check if the bone was exported
                                bone = self.get_bone( bonename )
                                _loc = bone.pose_location
                                _rot = bone.pose_rotation
                                _scl = bone.pose_scale

                                keyframe = doc.createElement('keyframe')
                                keyframe.setAttribute('time', str((frame-strip.frame_start)/_fps))
                                _keyframes[ bonename ].appendChild( keyframe )
                                trans = doc.createElement('translate')
                                keyframe.appendChild( trans )
                                x,y,z = _loc
                                trans.setAttribute('x', '%6f' %x)
                                trans.setAttribute('y', '%6f' %y)
                                trans.setAttribute('z', '%6f' %z)

                                rot =  doc.createElement( 'rotate' )
                                keyframe.appendChild( rot )
                                q = _rot
                                rot.setAttribute('angle', '%6f' %q.angle )
                                axis = doc.createElement('axis'); rot.appendChild( axis )
                                x,y,z = q.axis
                                axis.setAttribute('x', '%6f' %x )
                                axis.setAttribute('y', '%6f' %y )
                                axis.setAttribute('z', '%6f' %z )

                                scale = doc.createElement('scale')
                                keyframe.appendChild( scale )
                                x,y,z = _scl
                                scale.setAttribute('x', '%6f' %x)
                                scale.setAttribute('y', '%6f' %y)
                                scale.setAttribute('z', '%6f' %z)

                    if getConfig()['INDEPENDENT_ANIM']: strip.mute = True

                if getConfig()['INDEPENDENT_ANIM']:
                    nla.mute = True
                    for strip in nla.strips:
                        strip.mute = False

            if getConfig()['INDEPENDENT_ANIM']:
                for nla in arm.animation_data.nla_tracks:
                    nla.mute = False

        return doc.toprettyxml()
